from cProfile import label
import numpy as np
import os
from apk_process import process_apk


# apk_path_list = ["./sample_apk/000db93ae6b67e6303e5b5b4839df016.apk", 
#             "./sample_apk/000f3a904966dcffed17a965c183ba70.apk",
#             "./sample_apk/000316cad2ed6477e73d839f60115e5f.apk"]
apk_path_list = []
apk_test_or_train = []
apk_labels = []

for (dirpath, dirnames, filenames) in os.walk("./sample/test"):
    for filename in filenames:
        apkname = os.path.join(dirpath, filename)
        apk_istrain = 0
        if apkname.split('/')[-2] == "adware":
            apk_label = [1, 0, 0]
        elif apkname.split('/')[-2] == "normal":
            apk_label = [0, 1, 0]
        elif apkname.split('/')[-2] == "trojan":
            apk_label = [0, 0, 1]
        apk_path_list.append(apkname)
        apk_test_or_train.append(apk_istrain)
        apk_labels.append(apk_label)

for (dirpath, dirnames, filenames) in os.walk("./sample/train"):
    for filename in filenames:
        apkname = os.path.join(dirpath, filename)
        apk_istrain = 1
        if apkname.split('/')[-2] == "adware":
            apk_label = [1, 0, 0]
        elif apkname.split('/')[-2] == "normal":
            apk_label = [0, 1, 0]
        elif apkname.split('/')[-2] == "trojan":
            apk_label = [0, 0, 1]
        apk_path_list.append(apkname)
        apk_test_or_train.append(apk_istrain)
        apk_labels.append(apk_label)

# apk_path_list = apk_path_list[0:15]
# apk_test_or_train = apk_test_or_train[0:15]
# apk_labels = apk_labels[0:15]


per_dict = {}  # permission
cla_dict = {}  # class
lib_dict = {}  # library
int_dict = {}  # interface
met_dict = {}  # method

'''
class Myapk():
    def __init__(self, permission, apk_class, library, interface, methods):
        self.permission = permission
        self.apk_class = apk_class
        self.library = library
        self.interface = interface
        self.methods = methods
        self.per_vector = []
        self.cla_vector = []
        self.lib_vector = []
        self.int_vector = []
        self.met_vector = [] 
'''

# per_feature = []    # permission
# cla_feature = []    # class
# lib_feature = []    # library
# int_feature = []    # interface
# met_feature = []    # method

# Collect all features
apk_list = []
apk_test_or_train_new = []
apk_labels_new = []
apk_path_list_new = []
count = 0
for apk_idx in range(len(apk_path_list)):
    try:
        apk = process_apk(apk_path_list[apk_idx], per_dict, cla_dict, lib_dict, int_dict, met_dict)   
        # per_feature.extend(apk.permission)
        # per_feature = list(set(per_feature))
        
        # cla_feature.extend(apk.apk_class)
        # cla_feature = list(set(cla_feature))
        
        # lib_feature.extend(apk.library)
        # lib_feature = list(set(lib_feature))
        
        # int_feature.extend(apk.interface)
        # int_feature = list(set(int_feature))
        
        # met_feature.extend(apk.methods)
        # met_feature = list(set(met_feature))
        
        apk_list.append(apk)
        apk_test_or_train_new.append(apk_test_or_train[apk_idx])
        apk_labels_new.append(apk_labels[apk_idx])
        apk_path_list_new.append(apk_path_list[apk_idx])
        count += 1
        print(apk_path_list[apk_idx] + " Done! " + str(count))
    except:
        print(apk_path_list[apk_idx] + " failed!")
        continue
    # if count == 100:
    #     break
    
    
    
apk_test_or_train_new = np.array(apk_test_or_train_new)
apk_labels_new = np.array(apk_labels_new)
    
# judge 0 or 1 
def fun(feature, self_feature_list):
    if feature in self_feature_list:
        return 1
    else:
        return 0 

# Construct vector for every APK
apk_feature_list = []
count = 0
for apk in apk_list:
    apk_feature = []
    ''' This part optimized
    apk.per_vector = np.array([fun(x, apk.permission) for x in list(per_dict.values())])
    apk.cla_vector = np.array([fun(x, apk.apk_class) for x in list(cla_dict.values())])
    apk.lib_vector = np.array([fun(x, apk.library) for x in list(lib_dict.values())])
    apk.int_vector = np.array([fun(x, apk.interface) for x in list(int_dict.values())])
    apk.met_vector = np.array([fun(x, apk.methods) for x in list(met_dict.values())])
    '''
    
    apk.per_vector = [0] * len(per_dict.values())
    for item in apk.permission:
        apk.per_vector[item] = 1
    
    apk.cla_vector = [0] * len(cla_dict.values())
    for item in apk.apk_class:
        apk.cla_vector[item] = 1
        
    apk.lib_vector = [0] * len(lib_dict.values())
    for item in apk.library:
        apk.lib_vector[item] = 1
        
    apk.int_vector = [0] * len(int_dict.values())
    for item in apk.interface:
        apk.int_vector[item] = 1
        
    apk.met_vector = [0] * len(met_dict.values())
    for item in apk.methods:
        apk.met_vector[item] = 1

    apk.per_vector = np.array(apk.per_vector, dtype=bool)
    apk.cla_vector = np.array(apk.cla_vector, dtype=bool)
    apk.lib_vector = np.array(apk.lib_vector, dtype=bool)
    apk.int_vector = np.array(apk.int_vector, dtype=bool)
    apk.met_vector = np.array(apk.met_vector, dtype=bool)
        
    apk_feature.append(apk.per_vector)
    apk_feature.append(apk.cla_vector)
    apk_feature.append(apk.lib_vector)
    apk_feature.append(apk.int_vector)
    apk_feature.append(apk.met_vector)
    
    apk_feature_list.append(apk_feature)
    count += 1
    print(str(count) + "apk vector build")

apk_feature_list = np.array(apk_feature_list, dtype=object)

# Construct graph
permission_graph = np.zeros((len(apk_list), len(apk_list)), dtype=bool)
class_graph = np.zeros((len(apk_list), len(apk_list)), dtype=bool)
library_graph = np.zeros((len(apk_list), len(apk_list)), dtype=bool)
interface_graph = np.zeros((len(apk_list), len(apk_list)), dtype=bool)
method_graph = np.zeros((len(apk_list), len(apk_list)), dtype=bool)

count = 0
for i in range(len(apk_list)):
    for j in range(len(apk_list)):
        multi_result = np.multiply(apk_list[i].per_vector, apk_list[j].per_vector)
        if (1 in multi_result) or i==j:
            permission_graph[i][j] = 1
            
        multi_result = np.multiply(apk_list[i].cla_vector, apk_list[j].cla_vector)
        if (1 in multi_result) or i==j:
            class_graph[i][j] = 1
            
        multi_result = np.multiply(apk_list[i].lib_vector, apk_list[j].lib_vector)
        if (1 in multi_result) or i==j:
            library_graph[i][j] = 1
            
        multi_result = np.multiply(apk_list[i].int_vector, apk_list[j].int_vector)
        if (1 in multi_result) or i==j:
            interface_graph[i][j] = 1
            
        multi_result = np.multiply(apk_list[i].met_vector, apk_list[j].met_vector)
        if (1 in multi_result) or i==j:
            method_graph[i][j] = 1
    count += 1
    print(str(count) + "apk graph build")
            

import scipy.io as sio
import scipy.sparse as sparse
permission_graph = sparse.csr_matrix(permission_graph | np.eye(len(apk_list), dtype=bool))
class_graph = sparse.csr_matrix(class_graph | np.eye(len(apk_list), dtype=bool))
library_graph = sparse.csr_matrix(library_graph | np.eye(len(apk_list), dtype=bool))
interface_graph = sparse.csr_matrix(interface_graph | np.eye(len(apk_list), dtype=bool))
method_graph = sparse.csr_matrix(method_graph | np.eye(len(apk_list), dtype=bool))
sio.savemat("./apk_3.mat",mdict={"permission_graph":permission_graph,
                                 "class_graph":class_graph,
                                 "library_graph":library_graph,
                                 "interface_graph":interface_graph,
                                 "method_graph":method_graph, 
                                 "apk_labels":apk_labels_new, 
                                 "apk_test_or_train": apk_test_or_train_new,
                                #  "per_dict": per_dict,
                                #  "cla_dict": cla_dict,
                                #  "lib_dict": lib_dict,
                                #  "int_dict": int_dict,
                                #  "met_dict": met_dict,
                                 "apk_path_list": apk_path_list_new,
                                 "apk_feature_list": apk_feature_list}, long_field_names=True)

np.save("per_dict.npy", per_dict)
np.save("cla_dict.npy", cla_dict)
np.save("lib_dict.npy", lib_dict)
np.save("int_dict.npy", int_dict)
np.save("met_dict.npy", met_dict)
    
